package com.dongzili.common.mybatisplus;

import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.ParameterUtils;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.plugins.inner.PaginationInnerInterceptor;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.sql.SQLException;
import java.util.List;

/**
 * 自定义分页查询插件,缓存分页查询结果,只有第一次查询会去查询数据总行数
 */
public class CustomizePaginationInterceptor extends PaginationInnerInterceptor {

    public boolean willDoQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        IPage<?> page = (IPage<?>)ParameterUtils.findPage(parameter).orElse(null);
        if (page != null && page.getSize() >= 0L && page.searchCount() && resultHandler == Executor.NO_RESULT_HANDLER) {

            String totalCacheKey = PageTotalCache.generateKey(ms.getId(), page.getRecords());
            if (page.getCurrent() > 1){
                Long cachedTotal = PageTotalCache.getTotal(totalCacheKey);
                if (cachedTotal != null) {
                    page.setTotal(cachedTotal); // 设置缓存的总数
                    return true;
                }
            }

            MappedStatement countMs = this.buildCountMappedStatement(ms, page.countId());
            BoundSql countSql;
            if (countMs != null) {
                countSql = countMs.getBoundSql(parameter);
            } else {
                countMs = this.buildAutoCountMappedStatement(ms);
                String countSqlStr = this.autoCountSql(page, boundSql.getSql());
                PluginUtils.MPBoundSql mpBoundSql = PluginUtils.mpBoundSql(boundSql);
                countSql = new BoundSql(countMs.getConfiguration(), countSqlStr, mpBoundSql.parameterMappings(), parameter);
                PluginUtils.setAdditionalParameter(countSql, mpBoundSql.additionalParameters());
            }

            CacheKey cacheKey = executor.createCacheKey(countMs, parameter, rowBounds, countSql);
            List<Object> result = executor.query(countMs, parameter, rowBounds, resultHandler, cacheKey, countSql);
            long total = 0L;
            if (CollectionUtils.isNotEmpty(result)) {
                Object o = result.getFirst();
                if (o != null) {
                    total = Long.parseLong(o.toString());
                }
            }

            page.setTotal(total);
            PageTotalCache.putTotal(totalCacheKey, total);

            return this.continuePage(page);
        } else {
            return true;
        }
    }
}